import argparse
import paramiko

config = {
    "hosts": [
        {
            "hostname": "hostname0",
            "bin": "/path/to/torch/bin/torchrun",
            "main_dir": "/path/to/torlink",
            "log_dir": "/path/to/torlink/pylogs",
        },
        {
            "hostname": "hostname1",
            "bin": "/path/to/torch/bin/torchrun",
            "main_dir": "/path/to/torlink",
            "log_dir": "/path/to/torlink/pylogs",
        },
        {
            "hostname": "hostname2",
            "bin": "/path/to/torch/bin/torchrun",
            "main_dir": "/path/to/torlink",
            "log_dir": "/path/to/torlink/pylogs",
        },
        {
            "hostname": "hostname3",
            "bin": "/path/to/torch/bin/torchrun",
            "main_dir": "/path/to/torlink",
            "log_dir": "/path/to/torlink/pylogs",
        },
    ]
}


def arg_parser():
    parser = argparse.ArgumentParser(description="torch arg parser")
    parser.add_argument("-b", "--bootstrap_servers", type=str)
    parser.add_argument("-s", "--batch_size", type=int)
    parser.add_argument("-w", "--window_size", type=int, default=1024)
    parser.add_argument("-n", "--parallelism", type=int)
    parser.add_argument("-k", "--ck_interval", type=int)
    parser.add_argument("-x", "--topic-prefix", type=str, default="testtopic")
    parser.add_argument("-m", "--application", type=str)
    return parser


def run():
    args = arg_parser().parse_args()
    bootstrap_servers = args.bootstrap_servers
    batch_size = args.batch_size
    window_size = args.window_size
    parallelism = args.parallelism
    ck_interval = args.ck_interval
    topic_prefix = args.topic_prefix
    application = args.application
    topic = f"{topic_prefix}{parallelism}"

    main_file = f"{application}.py"

    procs_per_node = int(parallelism / len(config["hosts"]))

    for conf in config["hosts"]:
        hostname = conf["hostname"]
        bin = conf["bin"]
        main_dir = conf["main_dir"]
        log_dir = conf["log_dir"]
        main_py = f"{main_dir}/{main_file}"
        start(
            hostname,
            bin,
            main_py,
            bootstrap_servers,
            topic,
            batch_size,
            window_size,
            parallelism,
            procs_per_node,
            ck_interval,
            log_dir,
        )


def start(
    hostname,
    bin,
    main_py,
    bootstrap_servers,
    topic,
    batch_size,
    window_size,
    parallelism,
    procs_per_node,
    ck_interval,
    log_dir,
):
    pkey = paramiko.RSAKey.from_private_key_file("/path/to/key")
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(hostname=hostname, port=22, username="xxx", pkey=pkey)
    cmd = f"nohup {bin} --nnodes=4 --nproc_per_node={procs_per_node} --rdzv_id=100 --rdzv_backend=c10d --rdzv_endpoint=master:29401 {main_py} -b {bootstrap_servers} -t {topic} -s {batch_size} -p {parallelism} -k {ck_interval} -w {window_size} > {log_dir}/{main_py}_train.log 2>&1 &"
    stdin, stdout, stderr = ssh.exec_command(cmd)


if __name__ == "__main__":
    run()
